import subprocess
import time

from constant import *
from utils import logger
from env import env

from tool_agent.tool_agent import tool_agent


# 知识库 ： 查规则怎么写 -> 读文档
# 规则 ： 增删改查
# 输出 ： 总结输出

def action(action_func):
    """
    action的修饰器
    """

    def wrap(*kargs, **kwargs):
        logger.critical(
            f'Action <{action_func.__code__.co_name}> was called. Arguments: {kargs if kargs != () else ""}, '
            f'{kwargs if kwargs != {} else ""}')
        return action_func(*kargs, **kwargs)

    return wrap


def reload_rules():
    ret = subprocess.run("suricatasc -c reload-rules",
                         shell=True,
                         stdout=subprocess.PIPE)
    if 'OK' not in (out := ret.stdout.decode('utf-8')):
        logger.error(f"Reload rules failed, suricatasc response {out}")
    suricata_response = env.get_buffer()
    logger.critical(suricata_response)
    return suricata_response


def read_rules():
    with open(SURICATA_RULE_FILE, 'r') as f:
        rules = f.readlines()
    return rules


def write_rules(rules):
    with open(SURICATA_RULE_FILE, 'w') as f:
        f.writelines(rules)


def get_current_rules():
    with open(SURICATA_RULE_FILE, 'r') as f:
        lines = f.readlines()
    rules = "".join([
        f"{index}\t{line}" if line.strip() != '' else ''
        for index, line in enumerate(lines)
    ])
    return f"当前的入侵防御规则：\n{rules}"


@action
def get_document():
    pass


@action
def get_rules():
    return get_current_rules()


@action
def add_rule(rule_intent: str, lineno=None):
    rules = tool_agent.invoke(rule_intent)
    # 添加规则
    if lineno is None:
        with open(SURICATA_RULE_FILE, 'a') as f:
            for rule in rules:
                f.write(rule.strip() + '\n')
        # 更新规则
        ret = reload_rules()
        return f"新规则`{rules}`已添加，重载规则执行完毕,{get_current_rules()}\nSuricata response {ret}"
    else:
        with open(SURICATA_RULE_FILE, 'r+') as f:
            current_rules = f.readlines()
            current_line = lineno
            for rule in rules:
                current_rules.insert(current_line, rule + '\n')
                current_line += 1
            f.seek(0)
            f.writelines(current_rules)
        return f"新规则`{rules}`已添加于第{lineno}行，重载规则执行完毕,{get_current_rules()}"


@action
def del_rule(lineno):
    rules = read_rules()
    if lineno >= len(rules):
        return f"规则删除失败，不存在的行号{lineno};行号最大为{len(rules)}"
    delete_rule = rules[lineno]
    del rules[lineno]
    write_rules(rules)
    ret = reload_rules()
    return f"规则`{delete_rule}`已删除，重载规则执行完毕，{get_current_rules()}\n Suricata response {ret}"


@action
def change_rule(rule_intent, lineno):
    current_rules = read_rules()
    if lineno >= len(current_rules):
        return f"规则更新失败，不存在的行号{lineno};行号最大为{len(current_rules)}"

    rules = tool_agent.invoke(rule_intent)
    current_line = lineno
    for rule in rules:
        current_rules.insert(current_line, rule + '\n')
        current_line += 1
    write_rules(current_rules)
    ret = reload_rules()
    return f"规则已成功更新，{get_current_rules()}\nSuricata response {ret}"


@action
def get_log():
    with open(SURICATA_OUTPUT_FILE, 'r') as f:
        logs = f.read()
    return f"Suricata日志为:{logs}"


if __name__ == '__main__':
    # time.sleep(5)
    # add_rule('drop http2 any any -> any any (msg:"Alarm detected"; sid:1; rev:1;)')
    # time.sleep(5)
    # del_rule(0)
    print(get_log())
